#include "device/vector_path_v2.h"
#include <memory>
#define WidthSize 1024
#define BufferSize WidthSize * WidthSize

VectorPath_v2::VectorPath_v2(std::shared_ptr<MemInterface> In, std::shared_ptr<MemInterface> Out) {
  InMemPtr = std::move(In);
  OutMemPtr = std::move(Out);
  vector_start = false;
  irq_status_reg = false;
  irq_mask_reg = false;
  ofmBuffer = new s8[BufferSize];
  //ifm_w_step = 512;
  sram_addr_step = 16;
}

VectorPath_v2::~VectorPath_v2() {
  delete[] ofmBuffer;
}

inline s8 get_input(s8 *ifm_start, u32 index_w, u32 index_h, u32 ifm_w,
                     u32 ifm_h, u32 pad_left, u32 pad_right, u32 pad_top,
                     u32 pad_bottom, s8 pad_word, u32 ifm_w_step) {
  CHECK(index_w < ifm_w + pad_left + pad_right);
  CHECK(index_h < ifm_h + pad_top + pad_bottom);
  if ((index_w < pad_left) || (index_w >= ifm_w + pad_left) ||
      (index_h < pad_top) || (index_h >= ifm_h + pad_top))
    return pad_word;
  else {
    return ifm_start[ifm_w_step * (index_h - pad_top) + index_w - pad_left];
  }
}

void VectorPath_v2::element_wise(s8* ifm1, s8* ifm2) {
  for (u32 i = 0; i < ifm_size; i++) {
    s64 result = 0;
    if (vector_type == 1){
      u32 scale_shift1 = scale_shift & 0x1f;
      u32 scale_shift2 = (scale_shift & 0x3e0) >> 5;
      u32 scale1 = scale_num & 0xffff;
      u32 scale2 = (scale_num & 0xffff0000) >> 16;
      s64 lresult,rresult;
      lresult = (s64)ifm1[i] * (s64)scale1;
      rresult = (s64)ifm2[i] * (s64)scale2;
      if(scale_shift1 != 0)
        lresult = (lresult + ((s64)1 << (scale_shift1 - 1))) >> scale_shift1;
      if(scale_shift2 != 0)
        rresult = (rresult + ((s64)1 << (scale_shift2 - 1))) >> scale_shift2;
      result =  lresult + rresult;
    }
    else if (vector_type == 2){
      result = ((s64)ifm1[i]) * ((s64)ifm2[i]) * scale_num;
      u32 scale_shift1 = scale_shift & 0x3f;
      if(scale_shift1 != 0)
        result = (result+ ((s64)1 << (scale_shift1 - 1))) >> scale_shift1;
    }
    else
      CHECK_MSG(false, "error status\n");
    ofmBuffer[i] = sat_val(result);
  }
}

void VectorPath_v2::pooling(u64 ofm_addr, s8* ifm) {
  u32 pool_w = 0, pool_h = 0, pooling_type = 0;
  u32 ifm_w_step = 0;
  switch (vector_type) {
  case 1 << 2: // ofm = 2x2avepooling(ifm0)
    pool_w = 2;
    pool_h = 2;
    pooling_type = 0;
    break;
  case 1 << 3: // ofm = 3x3avepooling(ifm0)
    pool_w = 3;
    pool_h = 3;
    pooling_type = 0;
    break;
  case 1 << 4: // ofm = 4x4avepooling(ifm0)
    pool_w = 4;
    pool_h = 4;
    pooling_type = 0;
    break;
  case 1 << 5: // ofm = 2x2maxpooling(ifm0)
    pool_w = 2;
    pool_h = 2;
    pooling_type = 1;
    break;
  case 1 << 6: // ofm = 3x3maxpooling(ifm0)
    pool_w = 3;
    pool_h = 3;
    pooling_type = 1;
    break;
  case 1 << 7: // ofm = 4x4maxpooling(ifm0)
    pool_w = 4;
    pool_h = 4;
    pooling_type = 1;
    break;
  default:
    CHECK(false);
  }
  u32 ifm_width128byte = ifm_rd_step_config & 0xffff;
  u32 ifm_w_X_stride_h = (ifm_rd_step_config & 0xffff0000) >> 16;
  u32 stride_w = pooling_stride & 0xf;
  u32 stride_h = (pooling_stride & 0xf0) >> 4;
  u32 pad_left = padding_size & 0xf;
  u32 pad_right = (padding_size & 0xf0) >> 4;
  u32 pad_top = (padding_size & 0xf00) >> 8;
  u32 pad_bottom = (padding_size & 0xf000) >> 12;
  CHECK(ifm_width128byte * stride_h == ifm_w_X_stride_h);
  ifm_w_step = ifm_width128byte * 16;
  //CHECK((ifm_width + pad_left + pad_right - pool_w) / stride_w + 1 == ofm_width);
  //CHECK((ifm_height + pad_top + pad_bottom - pool_h) / stride_h + 1 == ofm_height);
  //u32 ifm_w_sum = pad_left + ifm_width + pad_right;
  //u32 ifm_h_sum = pad_top + ifm_height + pad_bottom;
  //make sure ifm_w/ifm_h < WidthSize
  //CHECK(ifm_w_sum < WidthSize);
  //CHECK(ifm_h_sum < WidthSize);
  //u64 input_addr = ifm_addr;
  u64 output_addr_tmp = ofm_addr;
  for(u32 ofm_h = 0;ofm_h < ofm_height;ofm_h++){
    for(u32 ofm_w =0;ofm_w < ofm_width;ofm_w++){
      u32 index_ifm_w_start = ofm_w * stride_w;
      u32 index_ifm_h_start = ofm_h * stride_h;
      if(pooling_type == 0){
        s64 result = 0;
        for(u32 knl_h = 0;knl_h < pool_h;knl_h++){
          for(u32 knl_w = 0;knl_w < pool_w;knl_w++){
            result += get_input(ifm, index_ifm_w_start + knl_w,
                                index_ifm_h_start + knl_h, ifm_width,
                                ifm_height, pad_left, pad_right, pad_top,
                                pad_bottom, padding_word, ifm_w_step);
          }
        }
        result = (s64)result * (s64)scale_num;
        u32 scale_shift1 = scale_shift & 0x3f;
        if(scale_shift1 != 0)
          result = (result + ((s64)1 << (scale_shift1 - 1))) >> scale_shift1;
        ofmBuffer[ofm_w] = sat_val(result);
      }
      else{
        s8 result = -128;
        for(u32 knl_h = 0;knl_h < pool_h;knl_h++){
          for(u32 knl_w = 0;knl_w < pool_w;knl_w++){
            s8 tmp = get_input(ifm, index_ifm_w_start + knl_w,
                               index_ifm_h_start + knl_h, ifm_width,
                               ifm_height, pad_left, pad_right, pad_top,
                               pad_bottom, padding_word, ifm_w_step);
            if(tmp > result) result = tmp;
          }
        }
        ofmBuffer[ofm_w] = result;
      }
      OutMemPtr->write_mem(output_addr_tmp, ofmBuffer, ofm_width);
    }
    output_addr_tmp += ofm_addr_w_offset;
  }
  return;
}

void VectorPath_v2::reduce_sum(s8* ifm) {
  switch (reduce_mode){
  case 1:
  {
    u32 ifm_step = ((ifm_width + sram_addr_step - 1) / sram_addr_step) * sram_addr_step;
    u32 ofm_length = ofm_width / 4;
    for(u32 index_h = 0;index_h < ofm_length;index_h++){
      s64 result = 0;
      for(u32 index_w = 0;index_w < ifm_width;index_w++){
        result += ifm[index_h * ifm_step + index_w];
      }
      result = (s64)result * (s64)scale_num;
      u32 scale_shift1 = scale_shift & 0x3f;
      if(scale_shift1 != 0)
        result = (result + ((s64)1 << (scale_shift1 - 1))) >> scale_shift1;
      ((s32*)ofmBuffer)[index_h] = (s32)result;
    }
    break;
  }
  case 2:
  {
    s64 result = 0;
    for(u32 index = 0;index < ifm_size;index++){
      result += ifm[index];
    }
    result = (s64)result * (s64)scale_num;
    u32 scale_shift1 = scale_shift & 0x3f;
    if(scale_shift1 != 0)
      result = (result + ((s64)1 << (scale_shift1 - 1))) >> scale_shift1;
    ((s32*)ofmBuffer)[0] = (s32)result;
    break;
  }
  default:
  {
    CHECK_MSG(false, "wrong reduce_mode\n");
  }
  }
  return;
}

void VectorPath_v2::reduce_mean(s8* ifm) {
  switch (reduce_mode){
  case 1:
  {
    u32 ifm_step = ((ifm_width + sram_addr_step - 1) / sram_addr_step) * sram_addr_step;
    u32 ofm_length = ofm_width;
    for(u32 index_h = 0;index_h < ofm_length;index_h++){
      s64 result = 0;
      for(u32 index_w = 0;index_w < ifm_width;index_w++){
        result += ifm[index_h * ifm_step + index_w];
      }
      result = (s64)result * (s64)scale_num;
      u32 scale_shift1 = scale_shift & 0x3f;
      if(scale_shift1 != 0)
        result = ((result + ((s64)1 << (scale_shift1 - 1))) >> scale_shift1);
      ofmBuffer[index_h] = sat_val(result);
    }
    break;
  }
  case 2:
  {
    s64 result = 0;
    for(u32 index = 0;index < ifm_size;index++){
      result += ifm[index];
    }
    result = (s64)result * (s64)scale_num;
    u32 scale_shift1 = scale_shift & 0x3f;
    if(scale_shift1 != 0)
      result = ((result + ((s64)1 << (scale_shift1 - 1))) >> scale_shift1);
    ofmBuffer[0] = sat_val(result);
    break;
  }
  default:
  {
    CHECK_MSG(false, "wrong reduce_mode\n");
  }
  }
  return;
}

void VectorPath_v2::max(s8* ifm) {
  switch (reduce_mode){
  case 0x1:
  {
    u32 ifm_step = ((ifm_width + sram_addr_step - 1) / sram_addr_step) * sram_addr_step;
    u32 ofm_length = ofm_width;
    for(u32 index_h = 0;index_h < ofm_length;index_h++){
      s8 result = -128;
      for(u32 index_w = 0;index_w < ifm_width;index_w++){
        s8 tmp = ifm[index_h * ifm_step + index_w];
        if(tmp > result) result = tmp;
      }
      ofmBuffer[index_h] = result;
    }
    break;
  }
  case 0x2:
  {
    s8 result = -128;
    for(u32 index = 0;index < ifm_size;index++){
      s8 tmp = ifm[index];
      if(tmp > result) result = tmp;
    }
    ofmBuffer[0] = result;
    break;
  }
  default:
  {
    CHECK_MSG(false, "wrong reduce_mode\n");
  }
  }
  return;
}

void VectorPath_v2::arg_max(s8* ifm) {
  switch (reduce_mode){
  case 0x1:
  {
    u32 ifm_step = ((ifm_width + sram_addr_step - 1) / sram_addr_step) * sram_addr_step;
    u32 ofm_length = ofm_width / 2;
    for(u32 index_h = 0;index_h < ofm_length;index_h++){
      u16 result = 0;
      s8 compare_val = -128;
      for(u32 index_w = 0;index_w < ifm_width;index_w++){
        s8 tmp = ifm[index_h * ifm_step + index_w];
        if(tmp > compare_val){
          compare_val = tmp;
          result = index_w;
        }
      }
      ((u16*)ofmBuffer)[index_h] = result;
    }
    break;
  }
  case 0x2:
  {
    u16 result = 0;
    s8 compare_val = -128;
    for(u32 index = 0;index < ifm_size;index++){
      s8 tmp = ifm[index];
      if(tmp > compare_val){
        compare_val = tmp;
        result = index;
      }
    }
    ((u16*)ofmBuffer)[0] = result;
    break;
  }
  default:
  {
    CHECK_MSG(false, "wrong reduce_mode\n");
  }
  }
  return;
}

s8 VectorPath_v2::sat_val(s64 result) {
  if (result > 127)
    return 127;
  if (result < -127)
    return -127;
  return (s8)result;
}

void VectorPath_v2::Run() {
  if (!vector_start)
    return;
  if (irq_status_reg)
    CHECK_MSG(false, "irq_status error");
  s8* ifm_addr[2];
  u64 ofm_addr;
  ifm_addr[0] = reinterpret_cast<s8*>(InMemPtr->get_ram_pointer()) + ((ifm_start_addr & 0xffff) << 4);
  ifm_addr[1] = reinterpret_cast<s8*>(InMemPtr->get_ram_pointer()) + (((ifm_start_addr & 0xffff0000) >> 16) << 4);
  ofm_addr = ((u64)ofm_start_h << 32) + ofm_start_l;
  switch (vector_type) {
  case 1:      // ofm = ifm0 + ifm1
  case 1 << 1: // ofm = ifm0 * ifm1
  {
    element_wise(ifm_addr[0], ifm_addr[1]);
    OutMemPtr->write_mem(ofm_addr, ofmBuffer, ifm_size);
    break;
  }
  case 1 << 2: // ofm = 2x2avepooling(ifm0)
  case 1 << 3: // ofm = 3x3avepooling(ifm0)
  case 1 << 4: // ofm = 4x4avepooling(ifm0)
  case 1 << 5: // ofm = 2x2maxpooling(ifm0)
  case 1 << 6: // ofm = 3x3maxpooling(ifm0)
  case 1 << 7: // ofm = 4x4maxpooling(ifm0)
  {
    pooling(ofm_addr, ifm_addr[0]);
    break;
  }
  case 1 << 8: // ofm = reduce_mean(ifm0)
  {
    reduce_mean(ifm_addr[0]);
    if(reduce_mode == 0x1){
      OutMemPtr->write_mem(ofm_addr, ofmBuffer, ofm_width);
    }
    else if(reduce_mode == 0x2){
      OutMemPtr->write_mem(ofm_addr, ofmBuffer, 1);
    }
    else{
      CHECK_MSG(false, "wrong reduce_mode!\n");
    }
    break;
  }
  case 1 << 9: // ofm = reduce_sum(ifm0)
  {
    reduce_sum(ifm_addr[0]);
    if(reduce_mode == 0x1){
      OutMemPtr->write_mem(ofm_addr, ofmBuffer, ofm_width);
    }
    else if(reduce_mode == 0x2){
      OutMemPtr->write_mem(ofm_addr, ofmBuffer, 4);
    }
    else{
      CHECK_MSG(false, "wrong reduce_mode!\n");
    }
    break;
  }
  case 1 << 10: // ofm = max(ifm0)
  {
    max(ifm_addr[0]);
    if(reduce_mode == 0x1){
      OutMemPtr->write_mem(ofm_addr, ofmBuffer, ofm_width);
    }
    else if(reduce_mode == 0x2){
      OutMemPtr->write_mem(ofm_addr, ofmBuffer, 1);
    }
    else{
      CHECK_MSG(false, "wrong reduce_mode!\n");
    }
    break;
  }
  case 1 << 11: // ofm = arg_max(ifm0)
  {
    arg_max(ifm_addr[0]);
    if(reduce_mode == 0x1){
      OutMemPtr->write_mem(ofm_addr, ofmBuffer, ofm_width);
    }
    else if(reduce_mode == 0x2){
      OutMemPtr->write_mem(ofm_addr, ofmBuffer, 2);
    }
    else{
      CHECK_MSG(false, "wrong reduce_mode!\n");
    }
    break;
  }
  default:
    CHECK(false);
  }
  vector_start = false;
  irq_status_reg = true;
  return;
}
